TVP: TV is Almost All You Need for Speculative Decoding

I wanted to understand why certain loss functions work better than others for speculative decoding training.
On my side, I had a specific recipe I followed, and when I was unable to beat it with recipes from papers, I was certainly confused. Made me analyze and question what really makes my recipe work well right now.
I first asked why are so many people using surrogates of expected acceptance length. I realized that this is quite problematic, because of prefix attenuation during training. However, these surrogates themselves optimize for different acceptance settings and ended up using cross-entropy or KL loss bases. I realized that KL and CE don't exactly optimize for acceptance length under speculative sampling. What's more is that the baselines in literature used thinking=off and T=0, both settings that larger frontier models seldom, if ever, serve in production (when's the last time your ChatGPT didn't think).
I began now with expected acceptance length directly, but I derived a surrogate using Jensen's inequality applied to it directly, rather than using any kinds of heuristics. I then got a somewhat loss that could take really large values. This was TVP 0. I realized I needed to soften it somehow, so I generalized it by assuming a discount factor on the weights of later tokens that depended on the acceptance probability of the prefix. I discovered that this has a unique form, and it led me to TVP.
Finally, stuff clicked. I realized what regimes TV was better in and when KL was better. It made me realize I should anneal the TVP beta parameter to switch from coverage to acceptance, leading to the best result and became my paper's main method which is TVP anneal
What was most striking to me was how good TV was. TV basically beat or tied most of the published losses in the thinking T=1 setting. Because TV was so strong, that prompted me to name the paper "TV is almost all you need", since of course, TV still has limitations.
I ended up discovering much better recipes than public recipes for popular libraries like TorchSpec or SpecForge, whose training recipes I significantly improved.
I especially had a lot of fun doing the math for this paper! It took a lot of hard-work!!

• Read on arXiv

I hope you enjoyed reading the paper!!

Any feedback is greatly appreciated!